from doas.io import UnableToLoad, register
from doas.spec import Spectrum
import numpy
import datetime


class SpectraSuitePlugin:
    
    def __init__(self):
        self.name = "SpectraSuite Tab-delimited"
    
    def test(self, fp):
        
        try:
            if fp.readline().startswith("SpectraSuite Data File"):
                return True
            else:
                return False
        
        except:
            return False
    
    
    #def load(self, fp):
    def load_header(self,fp):
        
        #skip the first two lines
        fp.readline()
        fp.readline()
        
        #read the header data
        header = {}
        while True:
            line = fp.readline()
            if line.startswith('>'):
                break
            words = line.split(':', 1)
            header[words[0].strip()] = words[1].strip()
        
        #TODO - this should return a SpectrumHeader object
        return header    

        
    def load(self, fp):
        header = self.load_header(fp)
        n_pix = int(header['Number of Pixels in Processed Spectrum'])        

        #read the actual spectral data
        data = numpy.fromfile(fp, count=2*n_pix, sep=" ")        
        wavelengths = data[::2]
        counts = data[1::2]
        
        #sanity check that we managed to load as much data as we wanted
        if data.shape[0] != 2*n_pix:
            raise UnableToLoad("Data loaded is wrong length")
        
        #convert the header data and create a Spectrum object
        capture_time = datetime.datetime.strptime(header['Date'],"%a %b %d %H:%M:%S %Z %Y")
        spectrometer_id = header['Spectrometers'].strip()
        int_time = float(header['Integration Time (usec)'].split()[0])/1e6
        n_avg = int(header['Spectra Averaged'].split()[0])
        
        return Spectrum(fp.name, wavelengths, counts, capture_time, spectrometer_id,
                        int_time, n_avg)
    
    
    def save(self, spectrum, filename):
        #TODO - error handling
        if not filename.endswith(".txt"):
            filename += '.txt'
            
        with open(filename,'w') as ofp:
            ofp.write('SpectraSuite Data File\n')
            ofp.write('++++++++++++++++++++++++++++++++++++\n')
            ofp.write('Date: ' + spectrum.capture_time.strftime("%a %b %d %H:%M:%S GMT %Y\n"))
            ofp.write('User: User\n')
            ofp.write('Dark Spectrum Present: No\n')
            ofp.write('Reference Spectrum Present: No\n')
            ofp.write('Number of Sampled Component Spectra: 1\n')
            ofp.write('Spectrometers: '+ spectrum.spectrometer_id + '\n')
            ofp.write('Integration Time (usec): ' + str(spectrum.int_time * 1e6) + ' ('+spectrum.spectrometer_id+')\n')
            ofp.write('Spectra Averaged: '+ str(int(spectrum.n_avg)) + ' ('+spectrum.spectrometer_id+')\n')
            ofp.write('Boxcar Smoothing: 0'+ ' ('+spectrum.spectrometer_id+')\n')
            ofp.write('Correct for Electrical Dark: No' + ' ('+spectrum.spectrometer_id+')\n')
            ofp.write('Strobe/Lamp Enabled: No'+ ' ('+spectrum.spectrometer_id+')\n')
            ofp.write('Correct for Detector Non-linearity: No' + ' ('+spectrum.spectrometer_id+')\n')
            ofp.write('Correct for Stray Light: No' + ' ('+spectrum.spectrometer_id+')\n')
            ofp.write('Number of Pixels in Processed Spectrum: '+str(len(spectrum.wavelengths))+'\n')
            ofp.write('>>>>>Begin Processed Spectral Data<<<<<\n')
            
            for i in range(len(spectrum.wavelengths)):
                ofp.write(str(spectrum.wavelengths[i])+'\t'+str(spectrum.counts[i])+'\n')
            
            ofp.write('>>>>>End Processed Spectral Data<<<<<\n')
            
            
    
register(SpectraSuitePlugin())        
